import operator, time
from typing import Annotated, Any
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

class State(TypedDict):
    aggregate: Annotated[list, operator.add]

class ReturnNodeValue:
    def __init__(self, node_secret: str):
        self._value = node_secret

    def __call__(self, state: State):
        print(f"Adding {self._value} to {state['aggregate']}")
        if self._value == 'c':
            time.sleep(10)
        return {"aggregate": [self._value]}

memory = MemorySaver()

builder = StateGraph(State)
builder.add_node("a", ReturnNodeValue("a"))
builder.add_node("b", ReturnNodeValue("b"))
builder.add_node("b2", ReturnNodeValue("b2"))
builder.add_node("c", ReturnNodeValue("c"))
builder.add_node("d", ReturnNodeValue("d"))

builder.add_edge(START, "a")
builder.add_edge("a", "b")
builder.add_edge("a", "c")
builder.add_edge("b", "b2")
# builder.add_edge("b2", "d")
# builder.add_edge("c","d")
builder.add_edge(["b2", "c"], "d")
builder.add_edge("d", END)
graph = builder.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "foo"}}
graph.invoke({"aggregate": []}, config)
print(graph.get_state(config))